// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/ir/_base/string_ref.h>
#include <tbir/runtime/object_internal.h>

#include <initializer_list>

#include <tbir/ir/_base/reflection.h>
#include <tbir/runtime/container/string_helper.h>
#include <tbir/runtime/object.h>
#include <tbir/runtime/registry.h>
#include <tbir/runtime/str_escape.h>

namespace tbir::runtime {
    // SEQualReduce traits for runtime containers.
    struct StringNodeTrait {
        static constexpr const std::nullptr_t VisitAttrs = nullptr;

        static void SHashReduce(const StringNode *key, SHashReducer hash_reduce) {
            hash_reduce->SHashReduceHashedValue(
                    BytesHash(key->data_container.data(), key->data_container.size()));
        }

        static bool SEqualReduce(const StringNode *lhs, const StringNode *rhs, SEqualReducer equal) {
            if (lhs == rhs)
                return true;
            return lhs->data_container == rhs->data_container;
        }
    };

    StringNode::StringNode() {
    }

    StringNode::StringNode(DataContainer o) : data_container(std::move(o)) {
    }

    StringNode::operator StringNode::self_view() const noexcept {
        return data_container.view();
    }

    TBIR_REGISTER_OBJECT_TYPE(StringNode);
    TBIR_REGISTER_REFLECTION_VTABLE(StringNode, StringNodeTrait)
            .set_creator([](const String &bytes) { return ObjectInternal::GetObjectPtr(StringRef(bytes)); })
            .set_repr_bytes([](const Object *n) -> String {
                return static_cast<const StringNode *>(n)->data_container;
            });

    TBIR_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
    .set_dispatch<StringNode>([](const ObjectRef &node, ReprPrinter *p) {
        auto *op = static_cast<const StringNode *>(node.get());
        p->stream << "b\'"
                  << runtime::BytesEscape(op->data_container.data(), op->data_container.size())
                  << "\'";
    });

    /******************************************************************************
     * StringRef container
     *****************************************************************************/

    const StringNode *StringRef::get() const {
        MX_DPTR(String);
        return d;
    }

    const StringRef::value_type *StringRef::data() const {
        MX_CHECK_DPTR(String);
        return d->data_container.data();
    }

    StringRef::operator StringRef::self_view() const noexcept {
        MX_DPTR(String);
        return d ? self_view{d->operator self_view()} : self_view();
    }

    StringRef::self_view StringRef::view() const noexcept {
        MX_DPTR(String);
        return d ? self_view{d->operator self_view()} : self_view();
    }

    StringRef::operator String() const {
        MX_DPTR(String);
        return d ? d->data_container : String{};
    }

    // default constructors
    StringRef::StringRef() {
        data_ = make_object<StringNode>();
    }

    //  constructors from other
    StringRef::StringRef(String other) {
        data_ = make_object<StringNode>(std::move(other));
    }

    StringRef::StringRef(const value_type *const data) : StringRef(String(data)) {
    }

    StringRef::StringRef(const value_type *data, size_type len) : StringRef(String(data, len)) {
    }

    // assign from other
    StringRef &StringRef::operator=(String other) {
        MX_DPTR(String);
        if (d) {
            d->data_container = std::move(other);
        } else {
            data_ = make_object<StringNode>(std::move(other));
        }
        return *this;
    }

    StringRef &StringRef::operator=(const value_type *other) {
        return operator=(String(other));
    }

    // Member functions
    StringNode *StringRef::CopyOnWrite() {
        if (data_.get() == nullptr) {
            data_ = make_object<StringNode>();
        } else if (!data_.unique()) {
            auto fbs = GetStringNode()->data_container;
            data_ = make_object<StringNode>(fbs);
        }
        return GetStringNode();
    }

    int StringRef::compare(const StringRef &other) const {
        return view().compare(other.view());
    }

    int StringRef::compare(const String &other) const {
        return view().compare(other);
    }

    int StringRef::compare(const char *other) const {
        return view().compare(other);
    }

    const char *StringRef::c_str() const {
        MX_CHECK_DPTR(String);
        return d->data_container.c_str();
    }

    int64_t StringRef::size() const {
        return view().size();
    }

    int64_t StringRef::length() const {
        return size();
    }

    bool StringRef::empty() const {
        return size() == 0;
    }

    StringRef StringRef::Concat(self_view lhs, self_view rhs) {
        StringNode::DataContainer container;
        container.reserve(lhs.size() + rhs.size());
        container.append(lhs.data(), lhs.size());
        container.append(rhs.data(), rhs.size());
        return StringRef(std::move(container));
    }

    // Private functions

    StringNode *StringRef::GetStringNode() const {
        return static_cast<StringNode *>(data_.get());
    }

    StringNode *StringRef::CreateOrGetStringNode() {
        if (!data_.get()) {
            data_ = make_object<StringNode>();
        }
        return static_cast<StringNode *>(data_.get());
    }

    TBIR_REGISTER_GLOBAL("runtime.String").set_body_typed([](String str) {
        return StringRef(std::move(str));
    });

    TBIR_REGISTER_GLOBAL("runtime.GetFFIString").set_body_typed([](StringRef str) {
        return str.operator String();
    });

    // runtime member function
    TBIR_REGISTER_GLOBAL("runtime.StringLen").set_body_typed([](StringRef str) {
        return static_cast<int64_t>(str.size());
    });

    TBIR_REGISTER_GLOBAL("runtime.StringAdd").set_body_typed([](StringRef lhs, StringRef rhs) {
        return lhs + rhs;
    });

    TBIR_REGISTER_GLOBAL("runtime.StringEqual").set_body_typed([](StringRef lhs, StringRef rhs) {
        return lhs == rhs;
    });

    TBIR_REGISTER_GLOBAL("runtime.StringHash").set_body_typed([](StringRef str) {
        return static_cast<int64_t>(std::hash<StringRef>()(str));
    });

    /******************************************************************************
     * String iterators
     *****************************************************************************/

    typename StringRef::iterator StringRef::begin() {
        auto n = CreateOrGetStringNode();
        return const_cast<StringRef::iterator>(n->data_container.data());
    }

    typename StringRef::const_iterator StringRef::begin() const {
        auto n = GetStringNode();
        MXCHECK(n != nullptr) << "[String.begin] container is null";
        return n->data_container.data();
    }

    typename StringRef::iterator StringRef::end() {
        auto n = CreateOrGetStringNode();
        return const_cast<StringRef::iterator>(n->data_container.data() + n->data_container.length());
    }

    typename StringRef::const_iterator StringRef::end() const {
        auto n = GetStringNode();
        MXCHECK(n != nullptr) << "[String.end] container is null";
        return n->data_container.data() + n->data_container.length();
    }

    typename StringRef::reverse_iterator StringRef::rbegin() {
        return reverse_iterator(end());
    }

    typename StringRef::const_reverse_iterator StringRef::rbegin() const {
        return const_reverse_iterator(end());
    }

    typename StringRef::reverse_iterator StringRef::rend() {
        return reverse_iterator(begin());
    }

    typename StringRef::const_reverse_iterator StringRef::rend() const {
        return const_reverse_iterator(begin());
    }

}  // namespace tbir::runtime
